import DataStructure.ProofData;
import DataStructure.hashMxGenData;
import algebra.curves.barreto_naehrig.*;
import algebra.curves.barreto_naehrig.abstract_bn_parameters.AbstractBNG1Parameters;
import algebra.curves.barreto_naehrig.abstract_bn_parameters.AbstractBNG2Parameters;
import algebra.curves.barreto_naehrig.abstract_bn_parameters.AbstractBNGTParameters;
import algebra.curves.barreto_naehrig.bn254a.BN254aFields;
import algebra.curves.barreto_naehrig.bn254a.BN254aG1;
import algebra.curves.barreto_naehrig.bn254a.BN254aG2;
import algebra.curves.barreto_naehrig.bn254a.BN254aPairing;
import algebra.curves.barreto_naehrig.bn254a.bn254a_parameters.BN254aG1Parameters;
import algebra.curves.barreto_naehrig.bn254a.bn254a_parameters.BN254aG2Parameters;
import com.google.common.hash.BloomFilter;
import com.google.common.hash.Funnel;
import com.google.common.hash.Funnels;
import com.google.zxing.common.reedsolomon.GenericGF;
import com.google.zxing.common.reedsolomon.ReedSolomonDecoder;
import com.google.zxing.common.reedsolomon.ReedSolomonEncoder;
import com.google.zxing.common.reedsolomon.ReedSolomonException;
import configuration.Configuration;
import it.unisa.dia.gas.jpbc.Element;
import it.unisa.dia.gas.jpbc.Pairing;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.storage.StorageLevel;
import profiler.generation.R1CSConstruction;
import profiler.utils.SparkUtils;
import relations.objects.Assignment;
import relations.r1cs.R1CSRelation;
import relations.r1cs.R1CSRelationRDD;
import scala.Tuple3;
import zk_proof_systems.zkSNARK.*;
import zk_proof_systems.zkSNARK.objects.CRS;
import zk_proof_systems.zkSNARK.objects.Proof;

import java.math.BigInteger;
import java.util.HashMap;
import java.util.Map;

public class Scheme2 {
    int chunkNum;
    int blockNum;
    Pairing pairing;
    Element _g;
    Element _h;
    Element _xi;
    Element epsilon;
    byte[][][] _Dx;
    Element[] phi;
    Element[] eta;
    Element[] _theta;
    Element _acc;
    BloomFilter<byte[]> bloomFilter;

    public void setup(byte[][][] Mx){
        int row = Mx.length;
        int collumn = Mx[0].length;
        this.chunkNum = row;
        this.blockNum = collumn;
        this.pairing = Arithmetic.pairing;
        this._g = pairing.getG1().newElement().setToRandom().getImmutable();
        this._h = pairing.getG2().newElement().setToRandom().getImmutable();
        this.epsilon = pairing.getZr().newElement().setToRandom().getImmutable();
        this._xi = pairing.getZr().newElement().setToRandom().getImmutable();
        this._Dx = Mx;
        this.phi = new Element[collumn];
        for(int i = 0; i < collumn; i++){
            Element newi = pairing.getZr().newElement().set(i);
            Element phi_i = pairing.getG1().newElement().set(_g).powZn(pairing.getZr().newElement().set(_xi).powZn(newi)).getImmutable();
            phi[i] = phi_i;
        }
        this.eta = new Element[row];
        for(int i = 0; i < row; i++){
            Element newi = pairing.getZr().newElement().set(i);
            Element epsilon_i = pairing.getG1().newElement().set(_h).powZn(pairing.getZr().newElement().set(epsilon).powZn(newi)).getImmutable();
            eta[i] = epsilon_i;
        }
        this._theta = new Element[chunkNum];
        for(int i = 0; i < chunkNum; i++){
            this._theta[i] = pairing.getZr().newElement().setToZero().getImmutable();
        }
    }

    public byte[][][] encode(byte[][][] Mx){
        GenericGF gf = new GenericGF(285, 256, 0);
        ReedSolomonEncoder encoder = new ReedSolomonEncoder(gf);
        byte[][][] newMx = new byte[this.chunkNum][init.blockNum][init.bitNum + init.parityShards];
        int size = this.chunkNum * init.blockNum;
        bloomFilter = BloomFilter.create(Funnels.byteArrayFunnel(), size);
        for(int i = 0; i < Mx.length; i++){
            for(int j = 0; j < Mx[i].length; j++){
                //bloomfilter
                bloomFilter.put(Mx[i][j]);
                //ReedSolomon
                int[] intTemp = new int[init.bitNum + init.parityShards];
                //byte to int
                for(int m = 0; m < Mx[i][j].length; m++){
                    intTemp[m] = Mx[i][j][m] & 0xff;
                }
                encoder.encode(intTemp, init.parityShards);

                //int to byte
                for(int n = 0; n < intTemp.length; n++){
                    newMx[i][j][n] = (byte) intTemp[n];
                }
            }
        }
        return newMx;
    }


    public void tag_Gen(){
        byte[][][] Mx = this._Dx;
        Element[][] HMx = new Element[this.chunkNum][this.blockNum];
        for(int i = 0; i < this.chunkNum; i++){
            for(int j = 0; j < this.blockNum; j++){
                //use hash and set it to Zr
                HMx[i][j] = pairing.getZr().newElement().setFromHash(Mx[i][j],0, Mx[i][j].length).getImmutable();
            }
        }
        for(int i = 0; i < Mx.length; i++){
            Element Lixi = Arithmetic.pbcPolynomialiEval(HMx, this._xi, i).getImmutable();
            Element g_Lixi = pairing.getG1().newElement().set(this._g).powZn(Lixi).getImmutable();
            this._theta[i] = g_Lixi;
        }
        Element prod = pairing.getZr().newElement().setToOne().getImmutable();
        for(int i = 0; i < this._theta.length; i++){
            Element ZHtheta = pairing.getZr().newElement().setFromHash(this._theta[i].toBytes(),0, this._theta[i].getLengthInBytes()).getImmutable();
            Element ZHtheta_epsilon = pairing.getZr().newElement().set(ZHtheta).add(this.epsilon).getImmutable();
            prod = prod.mul(ZHtheta_epsilon);
        }
        this._acc = pairing.getG2().newElement().set(this._h).powZn(prod);
    }

    public boolean read(Map<Integer, Element> chal) throws ReedSolomonException {
        //decode
        GenericGF gf = new GenericGF(285, 256, 0);
        ReedSolomonDecoder decoder = new ReedSolomonDecoder(gf);
        byte[][][] Mx = this._Dx;
        byte[][][] realMx = new byte[this.chunkNum][init.blockNum][init.bitNum];
        for(int i : chal.keySet()){
            for(int j = 0; j < Mx[i].length;j++){
                int[] intTemp = new int[init.bitNum + init.parityShards];
                byte[] realBlock = new byte[init.bitNum];
                for(int k = 0; k < Mx[i][j].length; k++){
                    intTemp[k] = Mx[i][j][k] & 0xff;
                }
                decoder.decode(intTemp, init.parityShards);
                for(int m = 0; m < realBlock.length; m++){
                    realBlock[m] = (byte)intTemp[m];
                }
                realMx[i][j] = realBlock;
                if(!bloomFilter.mightContain(realMx[i][j])){
                    System.out.println("wrong!");
                    return false;
                }
            }
        }
        return true;
    }

    public Map<Integer, Element> challenge(int CNum){
        if(CNum < 0 || CNum > this.chunkNum + 1){
            System.out.println("CNum wrong!");
        }
        int[] selectedBlockSet = Arithmetic.getUniqueRandomNum(this.chunkNum, CNum);
        Map<Integer, Element> chal = new HashMap<>();
        for(int i:selectedBlockSet){
            chal.put(i,pairing.getZr().newElement().setToRandom().getImmutable());
        }
        return chal;
    }

    public String sha256(byte[] data){
        byte[] res = DigestUtils.sha256(data);
        BigInteger hashInt = new BigInteger(1, res);
        return hashInt.toString();
    }

    public hashMxGenData hashMxGen(Map<Integer, Element> chal){
        byte[][][] hashMx = this._Dx;
        Element[][] HMx = new Element[this.chunkNum][this.blockNum];
        Element HMx_all = pairing.getZr().newElement().setToZero();
        Element temp = pairing.getZr().newElement().setToZero();
//        BigInteger[][] HMx = new BigInteger[this.chunkNum][this.blockNum];
//        BigInteger HMx_all = BigInteger.ZERO;
//       BigInteger temp = BigInteger.ZERO;
        for(Integer i : chal.keySet()){
            for(int j = 0; j < hashMx[i].length; j++){
                Element hashVal = pairing.getZr().newElement().setFromHash(hashMx[i][j], 0, hashMx[i][j].length);
                Element juhe = pairing.getZr().newElement().set(chal.get(i)).mul(hashVal);
                HMx[i][j] = juhe;
                HMx_all.add(juhe);
                temp.add(juhe);
//                String hashVal = sha256(hashMx[i][j]);
//                BigInteger BigHashVal = new BigInteger(hashVal);
//                Element juhe = pairing.getZr().newElement().set(chal.get(i)).mul(BigHashVal);
//                HMx[i][j] = juhe;
//                HMx_all = HMx_all.add(juhe);
//                temp = temp.add(juhe);
            }

        }
        return new hashMxGenData(HMx, HMx_all, chunkNum, blockNum);
    }

    public ProofData proofGen(Map<Integer, Element> chal,Element[][] HMx){
        Element[] Lx_sum = new Element[this.blockNum];
        for(int i = 0; i < this.blockNum; i++){
            Lx_sum[i] = pairing.getZr().newElement().setToZero().getImmutable();
        }
        for(int i:chal.keySet()){
            for(int j = 0; j < this.blockNum; j++){
                Element newHMx = pairing.getZr().newElement().set(HMx[i][j]).getImmutable();
                Lx_sum[j] = Lx_sum[j].add(newHMx).getImmutable();
            }
        }
        Element wit = pairing.getG2().newElement().setToOne().getImmutable();
        Map<Integer, Element> theta_prf = new HashMap<>();
        Map<Integer, Element> dictD_C = new HashMap<>();
        for(int i = 0; i < this.chunkNum; i++){
            dictD_C.put(i, this._theta[i]);
        }
        for(int i:chal.keySet()){
            dictD_C.remove(i);
            theta_prf.put(i, this._theta[i]);
        }
        Element[] prodD_C = new Element[1];
        prodD_C[0] = pairing.getZr().newElement().setToOne();
        for(int i:dictD_C.keySet()){
            Element[] slice = new Element[2];
            slice[0] = pairing.getZr().newElement().setFromHash(dictD_C.get(i).toBytes(), 0, dictD_C.get(i).getLengthInBytes());
            slice[1] = pairing.getZr().newElement().setToOne();
            prodD_C = Arithmetic.polynomialMul(prodD_C, slice);
        }

        for(int i = 0; i < prodD_C.length; i++){
            wit = wit.mul(this.eta[i].powZn(prodD_C[i]));
        }
        return new ProofData(theta_prf, wit, Lx_sum);
    }


    public boolean verify(Map<Integer, Element> chal, Map<Integer, Element> theta_prf, Element wit, Element[] Lx_sum){
        Element left1 = pairing.pairing(this._acc, this._h).getImmutable();
        Element prod = pairing.getZr().newElement().setToOne().getImmutable();
        Element left2 = pairing.getG1().newElement().setToOne().getImmutable();
        for(int i: chal.keySet()){
            Element zthetai = pairing.getZr().newElement().setFromHash(this._theta[i].toBytes(), 0, this._theta[i].getLengthInBytes()).getImmutable();
            Element zthetai_epsilon = zthetai.add(epsilon);
            prod = prod.mul(zthetai_epsilon);
            left2 = left2.mul(pairing.getG1().newElement().set(theta_prf.get(i)).powZn(chal.get(i)));
        }
        Element e4 = pairing.getG2().newElement().set(this._h).powZn(prod);
        Element right1 = pairing.pairing(wit,e4);
        System.out.println("pair1 -> "+left1.isEqual(right1));
        Element right2 = pairing.getG1().newElement().setToOne().getImmutable();
        for(int i = 0; i < this.blockNum; i++){
            right2 = right2.mul(pairing.getG1().newElement().set(this.phi[i]).powZn(Lx_sum[i]));
        }
        Element px = Arithmetic.polynomialEval(Lx_sum, this._xi);
        Element g_px = pairing.getG1().newElement().set(this._g).powZn(px);

        //System.out.println("3:"+px.isEqual(Arithmetic.polynomialEval(Lx_sum, this._xi)));

        //System.out.println("g_px.isEqual(right2) "+g_px.isEqual(right2));

        return left2.isEqual(right2);
    }

    public Configuration setUp() {
        Configuration config;
        config = new Configuration();
        config.setRuntimeFlag(false);
        config.setDebugFlag(true);
        return config;
    }

    private transient JavaSparkContext sc;
    Configuration config;

    public void sparkSetup(){
        final SparkConf conf = new SparkConf().setMaster("local").setAppName("default");
        conf.set("spark.files.overwrite", "true");
        conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        conf.registerKryoClasses(SparkUtils.zksparkClasses());

        sc = new JavaSparkContext(conf);

        config = new Configuration(2, 1, 1, 2, sc, StorageLevel.MEMORY_ONLY());
        config.setRuntimeFlag(false);
        config.setDebugFlag(true);
    }


    public static <BNFrT extends BNFields.BNFr<BNFrT>,
            BNFqT extends BNFields.BNFq<BNFqT>,
            BNFq2T extends BNFields.BNFq2<BNFqT, BNFq2T>,
            BNFq6T extends BNFields.BNFq6<BNFqT, BNFq2T, BNFq6T>,
            BNFq12T extends BNFields.BNFq12<BNFqT, BNFq2T, BNFq6T, BNFq12T>,
            BNG1T extends BNG1<BNFrT, BNFqT, BNG1T, BNG1ParametersT>,
            BNG2T extends BNG2<BNFrT, BNFqT, BNFq2T, BNG2T, BNG2ParametersT>,
            BNGTT extends BNGT<BNFqT, BNFq2T, BNFq6T, BNFq12T, BNGTT, BNGTParametersT>,
            BNG1ParametersT extends AbstractBNG1Parameters<BNFrT, BNFqT, BNG1T, BNG1ParametersT>,
            BNG2ParametersT extends AbstractBNG2Parameters<BNFrT, BNFqT, BNFq2T, BNG2T, BNG2ParametersT>,
            BNGTParametersT extends AbstractBNGTParameters<BNFqT, BNFq2T, BNFq6T, BNFq12T, BNGTT, BNGTParametersT>,
            BNPublicParametersT extends BNPublicParameters<BNFqT, BNFq2T, BNFq6T, BNFq12T>,
            BNPairingT extends BNPairing<BNFrT, BNFqT, BNFq2T, BNFq6T, BNFq12T, BNG1T, BNG2T, BNGTT, BNG1ParametersT, BNG2ParametersT, BNGTParametersT, BNPublicParametersT>>
    void SerialBNProofSystem(
            Configuration config,
            BNFrT fieldFactory,
            BNG1T g1Factory,
            BNG2T g2Factory,
            BNPairingT pairing,
            byte[][][] Mx,
            Map<Integer, Element> chal) {
        final Tuple3<R1CSRelation<BNFrT>, Assignment<BNFrT>, Assignment<BNFrT>> construction =
                R1CSConstruction.newtest(fieldFactory, config, Mx, chal);
        final R1CSRelation<BNFrT> r1cs = construction._1();
        final Assignment<BNFrT> primary = construction._2();
        final Assignment<BNFrT> fullAssignment = construction._3();

        final CRS<BNFrT, BNG1T, BNG2T, BNGTT> CRS =
                SerialSetup.generate(r1cs, fieldFactory, g1Factory, g2Factory, pairing, config);

        final Proof<BNG1T, BNG2T> proof =
                SerialProver.prove(CRS.provingKey(), primary, fullAssignment, fieldFactory, config);

        final boolean isValid = Verifier.verify(CRS.verificationKey(), primary, proof, pairing, config);

        System.out.println(isValid);
    }

   public static <BNFrT extends BNFields.BNFr<BNFrT>,
            BNFqT extends BNFields.BNFq<BNFqT>,
            BNFq2T extends BNFields.BNFq2<BNFqT, BNFq2T>,
            BNFq6T extends BNFields.BNFq6<BNFqT, BNFq2T, BNFq6T>,
            BNFq12T extends BNFields.BNFq12<BNFqT, BNFq2T, BNFq6T, BNFq12T>,
            BNG1T extends BNG1<BNFrT, BNFqT, BNG1T, BNG1ParametersT>,
            BNG2T extends BNG2<BNFrT, BNFqT, BNFq2T, BNG2T, BNG2ParametersT>,
            BNGTT extends BNGT<BNFqT, BNFq2T, BNFq6T, BNFq12T, BNGTT, BNGTParametersT>,
            BNG1ParametersT extends AbstractBNG1Parameters<BNFrT, BNFqT, BNG1T, BNG1ParametersT>,
            BNG2ParametersT extends AbstractBNG2Parameters<BNFrT, BNFqT, BNFq2T, BNG2T, BNG2ParametersT>,
            BNGTParametersT extends AbstractBNGTParameters<BNFqT, BNFq2T, BNFq6T, BNFq12T, BNGTT, BNGTParametersT>,
            BNPublicParametersT extends BNPublicParameters<BNFqT, BNFq2T, BNFq6T, BNFq12T>,
            BNPairingT extends BNPairing<BNFrT, BNFqT, BNFq2T, BNFq6T, BNFq12T, BNG1T, BNG2T, BNGTT, BNG1ParametersT, BNG2ParametersT, BNGTParametersT, BNPublicParametersT>>
    void DistributedBNProofSystem(
            Configuration config,
            BNFrT fieldFactory,
            BNG1T g1Factory,
            BNG2T g2Factory,
            BNPairingT pairing,
            byte[][][] Mx,
            Map<Integer, Element> chal) {

        final Tuple3<R1CSRelationRDD<BNFrT>, Assignment<BNFrT>, JavaPairRDD<Long, BNFrT>> construction =
                R1CSConstruction.newtest2(fieldFactory, config, Mx, chal);

        final R1CSRelationRDD<BNFrT> r1cs = construction._1();
        final Assignment<BNFrT> primary = construction._2();
        final JavaPairRDD<Long, BNFrT> fullAssignment = construction._3();

        final CRS<BNFrT, BNG1T, BNG2T, BNGTT> CRS =
                DistributedSetup.generate(r1cs, fieldFactory, g1Factory, g2Factory, pairing, config);

       long t_zkProof = System.currentTimeMillis();
        final Proof<BNG1T, BNG2T> proof =
                DistributedProver.prove(CRS.provingKeyRDD(), primary, fullAssignment, fieldFactory, config);
        long t_zkProof_end = System.currentTimeMillis();
       System.out.println("zkProof time"+(double)(t_zkProof_end - t_zkProof)/1000+"s");

       long t_zkVerify = System.currentTimeMillis();
        final boolean isValid = Verifier.verify(CRS.verificationKey(), primary, proof, pairing, config);
        long t_zkVerify_end = System.currentTimeMillis();
       System.out.println("zkVerify time"+(double)(t_zkVerify_end - t_zkVerify)/1000+"s");

        System.out.println(isValid);
    }


    public void SerialBN254aProofSystem(Configuration config,  byte[][][] Mx, Map<Integer, Element> chal) {
        final BN254aFields.BN254aFr fieldFactory = BN254aFields.BN254aFr.ONE;
        final BN254aG1 g1Factory = BN254aG1Parameters.ONE;
        final BN254aG2 g2Factory = BN254aG2Parameters.ONE;
        final BN254aPairing pairing = new BN254aPairing();

        SerialBNProofSystem(config,fieldFactory, g1Factory, g2Factory, pairing, Mx, chal);
    }

    public void DistributedBN254aProofSystem(Configuration config,  byte[][][] Mx, Map<Integer, Element> chal) {
        final BN254aFields.BN254aFr fieldFactory = new BN254aFields.BN254aFr(1);
        final BN254aG1 g1Factory = BN254aG1Parameters.ONE;
        final BN254aG2 g2Factory = BN254aG2Parameters.ONE;
        final BN254aPairing pairing = new BN254aPairing();

        DistributedBNProofSystem(config, fieldFactory, g1Factory, g2Factory, pairing,Mx, chal);
    }

    public static void main(String[] args) throws ReedSolomonException {
        Scheme2 scheme = new Scheme2();
        long t_fileGen = System.currentTimeMillis();
        byte[][][] Mx = Arithmetic.hashFilePolyGen(init.chunckNum, 0, init.blockNum);
        long t_fileGen_end = System.currentTimeMillis();
        System.out.println("file_ploy_gen time:"+ (double)(t_fileGen_end-t_fileGen)/1000+"s");

        byte[][][] realMx = scheme.encode(Mx);
        long t_init = System.currentTimeMillis();
        scheme.setup(realMx);
        long t_init_end = System.currentTimeMillis();
        System.out.println("Initialize time:"+(double)(t_init_end-t_init)/1000+"s");

        long t_tag = System.currentTimeMillis();
        scheme.tag_Gen();
        long t_tag_end = System.currentTimeMillis();
        System.out.println("tagGen time:"+(double)(t_tag_end-t_tag)/1000+"s");

        long t_chal = System.currentTimeMillis();
        Map<Integer, Element> chal = scheme.challenge(init.chunckNum);
        long t_chal_end = System.currentTimeMillis();
        System.out.println("challenge time:"+(double)(t_chal_end-t_chal)/1000+"s");

        boolean dataVerify = scheme.read(chal);
        if(dataVerify){
            System.out.println("data verify true!");
        }
        long t_tagGen_HMx = System.currentTimeMillis();
        hashMxGenData hashMxGen_data = scheme.hashMxGen(chal);
        long t_tagGen_HMx_end = System.currentTimeMillis();
        System.out.println("tagGen HMx time:"+(double)(t_tagGen_HMx_end-t_tagGen_HMx)/1000+"s");

        long t_proof = System.currentTimeMillis();
        ProofData proofData = scheme.proofGen(chal, hashMxGen_data.getHMx());
        long t_proof_end = System.currentTimeMillis();
        System.out.println("ProofGen time:"+(double)(t_proof_end-t_proof)/1000+"s");

        long t_verify = System.currentTimeMillis();
        boolean res = scheme.verify(chal, proofData.getTheta_prf(), proofData.getWit(), proofData.getLx_sum());
        long t_verify_end = System.currentTimeMillis();
        System.out.println("Verify time:"+(double)(t_verify_end-t_verify)/1000+"s");

        System.out.println("verification result: "+res);

        scheme.sparkSetup();
        //Configuration config = scheme.setUp();
       // scheme.SerialBN254aProofSystem(config, Mx, chal);

        scheme.DistributedBN254aProofSystem(scheme.config, Mx, chal);
    }



}
